package com.example.module.ai.api;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zhenxin.yu
 * @description DeepSeek-ollama本地部署模型
 * @date 2025-02-09 14:09
 **/
@Service
public class DeepSeekOllamaAiApiImpl implements AiApi {

    @Resource(name = "deepseekChatClient")
    private ChatClient chatClient;

    @Override
    public String completion(String question) {
        Prompt prompt = new Prompt(new UserMessage(question));
        return chatClient.prompt(prompt).call().content();
    }

    @Override
    public Flux<String> streamCompletion(String message) {
        AtomicInteger previousIndex = new AtomicInteger();
        Prompt prompt = new Prompt(new UserMessage(message));
        return chatClient.prompt(prompt).stream().content().map(
                content -> {
                    String newLetters = content.substring(previousIndex.get());
                    previousIndex.set(content.length());
                    return newLetters;
                }
        );
    }
}